#include <iostream>

#include "nets.hpp"

#include <fstream>
#include <vector>

int main(int argc, char* argv[])
{
    std::vector<float> rangeX{0, 4};
    const int numPoints = 301;
    const float dx      = (rangeX[1] - rangeX[0]) / (numPoints - 1);
    torch::Tensor x     = torch::zeros({numPoints, 1});
    torch::Tensor y     = torch::zeros({numPoints, 1});
    for (int i = 0; i < numPoints; ++i)
    {
        const float x_ = rangeX[0] + i * dx;
        x[i][0]        = x_;
        if (x_ < 1.0)
        {
            y[i][0] = x_;
        }
        else if (x_ > 3.0)
        {
            y[i][0] = x_ - 4.0;
        }
        else
        {
            y[i][0] = 2.0 - x_;
        }
        const auto random = torch::randn({1}) * 0.1;
        y[i][0] += random.item<float>();
    }

    std::shared_ptr<Net> net = std::make_shared<Net>();

    auto optimizer = std::make_shared<torch::optim::SGD>(net->parameters(), 1.e-3);

    auto dataSet = CustomDtaset(x, y).map(torch::data::transforms::Stack<>());
    auto dataLoader =
        torch::data::make_data_loader<torch::data::samplers::RandomSampler>(
            std::move(dataSet), 32);

    float lossVal = 1.0;
    int epoch     = 0;
    std::cout << std::scientific;
    while (std::abs(lossVal) > 1.e-3)
    {
        lossVal = 0.0;
        net->train();
        for (auto& batch : *dataLoader)
        {
            torch::Tensor outD   = net->forward(batch.data);
            torch::Tensor target = batch.target;
            torch::Tensor lossD  = torch::mse_loss(outD, target);
            optimizer->zero_grad();
            lossD.backward();
            optimizer->step();

            //
            lossVal += lossD.item<float>();
        }

        epoch++;
        std::cout << "Epoch.Idx : " << epoch << ", Train Loss : " << lossVal
                  << std::endl;

        net->eval();
        std::ofstream os("data_from_train_" + std::to_string(epoch) + ".csv");
        torch::Tensor out = net->forward(x);
        os << "x,label,pred" << std::endl;
        for (int i = 0; i < numPoints; ++i)
        {
            os << x[i][0].item<float>() << ", " << y[i][0].item<float>() << ", "
               << out[i][0].item<float>() << std::endl;
        }
        os.close();

        if (epoch >= 500) break;
    }

    return 0;
}
